import gensim
from gensim.models import Word2Vec
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models

# 示例代码片段，已标记化（tokenized）
# 在实际应用中，这些代码片段应为从代码文件中解析出的标识符、关键字、操作符等
code_snippets = [
    ['for', 'i', 'in', 'range', '(', '10', ')', ':'],
    ['if', 'x', '>', '0', ':'],
    ['print', '(', '"Hello"', ')'],
    ['def', 'function', '(', 'param', ')', ':'],
    # 更多代码片段
]

# 训练Word2Vec模型
word2vec_model = Word2Vec(sentences=code_snippets, vector_size=100, window=5, min_count=1, workers=4)

# 查看词汇表中的部分词语
print(word2vec_model.wv.index_to_key[:10])


# 定义代码片段的最大长度
max_length = 100  # 假设每个片段的最大标记数为10
vector_size = word2vec_model.vector_size


def vectorize_code_snippet(code_snippet):
    # 初始化矩阵，形状为 (max_length, vector_size)
    snippet_matrix = np.zeros((max_length, vector_size))

    for i, token in enumerate(code_snippet):
        if i < max_length:
            # 获取词向量，如果词不在词汇表中则使用零向量
            snippet_matrix[i] = word2vec_model.wv[token] if token in word2vec_model.wv else np.zeros(vector_size)

    return snippet_matrix


# 将所有代码片段转换为向量矩阵
vectorized_snippets = np.array([vectorize_code_snippet(snippet) for snippet in code_snippets])

# 增加一个通道维度以适应CNN输入 (batch_size, max_length, vector_size, 1)
vectorized_snippets = np.expand_dims(vectorized_snippets, axis=-1)

print("Vectorized code snippet shape:", vectorized_snippets.shape)


# 构建CNN模型
def build_cnn_model(input_shape):
    model = models.Sequential()

    # 第一层卷积层和池化层
    model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=input_shape))
    model.add(layers.MaxPooling2D((2, 2)))

    # 第二层卷积层和池化层
    model.add(layers.Conv2D(64, (3, 3), activation='relu'))
    model.add(layers.MaxPooling2D((2, 2)))

    # 第三层卷积层和池化层
    model.add(layers.Conv2D(128, (3, 3), activation='relu'))
    model.add(layers.MaxPooling2D((2, 2)))

    # 展平层
    model.add(layers.Flatten())

    # 全连接层
    model.add(layers.Dense(128, activation='relu'))

    # 输出层
    model.add(layers.Dense(1, activation='sigmoid'))

    return model


# 初始化模型
input_shape = (max_length, vector_size, 1)
model = build_cnn_model(input_shape)

# 编译模型
model.compile(optimizer='adam',
              loss='binary_crossentropy',
              metrics=['accuracy'])

# 打印模型摘要
model.summary()

# 模拟标签数据，1表示木马代码，0表示正常代码
labels = np.array([0, 1, 0, 1])  # 这是示例数据，真实情况下应为实际的标签

# 训练模型
history = model.fit(vectorized_snippets, labels, epochs=10, validation_split=0.2)

# 保存模型
model.save('trojan_detector_cnn_word2vec.h5')
